{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d6b482ba",
   "metadata": {},
   "source": [
    "### 0. Recap GNN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a130b97",
   "metadata": {},
   "source": [
    "##### Encoder Decoder Framework\n",
    "\n",
    "<img src=\"pics/encode-decode.png\" alt=\"Encoder Decoder Structure\" title=\"Encoder-Decoder\" width=\"600\"/>\n",
    "\n",
    "$Encoder$: maps each node to a low-dimensional vector \n",
    "- Shallow Encoders : Simplest encoding approach: Encoder is just an embedding-lookup\n",
    "\n",
    "> - O(|V|) parameters are needed; \n",
    "> - Tranductive: cannot generate embeddings for nodes that are not seen during training\n",
    "> - Do not incorporate node features\n",
    "\n",
    "- Deep Graph Encoders\n",
    "<br />\n",
    "\n",
    "$Dncoder$: predict score based on embedding to match node similarity. If supervised:\n",
    "\n",
    "> - Node classification: Predict a type of a given node\n",
    "> - Link prediction: Predict whether two nodes are linked\n",
    "> - Community detection: Identify densely linked clusters of nodes\n",
    "> - Network similarity: How similar are two (sub)net\n",
    "\n",
    "##### MESSAGE PASSING NETWORKS\n",
    "\n",
    "<img src=\"pics/gnn1.png\" alt=\"gnn1\" title=\"gnn1\" width=\"600\"/>\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{h}_i^{(k)} = \\gamma^{(k)} \\left( \\mathbf{h}_i^{(k-1)}, \\square_{j \\in \\mathcal{N}(i)} \\, \\phi^{(k)}\\left(\\mathbf{h}_i^{(k-1)}, \\mathbf{h}_j^{(k-1)},\\mathbf{e}_{j,i}\\right) \\right),\n",
    "\\end{equation}\n",
    "\n",
    "where $\\square$ denotes a differentiable, permutation invariant function, e.g., sum, mean or max, and $\\gamma$ and $\\phi$ denote differentiable functions such as MLPs (Multi Layer Perceptrons). The user only has to define these functions.\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12262eaa",
   "metadata": {},
   "source": [
    "### 1. RGCN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aeab619",
   "metadata": {},
   "source": [
    "#### Encoder\n",
    "\n",
    "The relational graph convolutional operator from the \"ModelingRelational Data with Graph Convolutional Networks\"\n",
    "    <https://arxiv.org/abs/1703.06103>`_ paper\n",
    "\n",
    "\\begin{equation}\n",
    "h_i^{(l+1)} = \\sigma(\\sum_{r\\in\\mathcal{R}}\n",
    "       \\sum_{j\\in\\mathcal{N}^r(i)}\\frac{1}{c_{i,r}}W_r^{(l)}h_j^{(l)}+W_0^{(l)}h_i^{(l)})\n",
    "\\end{equation}\n",
    "\n",
    " where : <br/>\n",
    " > - $\\mathcal{N}^r(i)$ is the neighbor set of node :$i$ w.r.t. relation $r$. <br/>\n",
    " > -  $c_{i,r}$ is the normalizer equal to $|\\mathcal{N}^r(i)|$. <br/>\n",
    "  > - $\\sigma$ is an activation function. <br/>\n",
    "  > - $W_0$ is the self-loop weight.<br/>\n",
    "\n",
    "\n",
    "where: $\\mathcal{R}$ denotes the set of relations, *i.e.* edge types. \n",
    "\n",
    "\n",
    "**basis Decomposition(Optional）**\n",
    "\n",
    "The problem of applying the above equation directly is the rapid growth of the number of parameters, especially with highly multi-relational data. In order to reduce model parameter size and prevent overfitting, the original paper proposes to use basis decomposition.\n",
    "\n",
    "The basis regularization decomposes $W_r$ by\n",
    "\\begin{equation}\n",
    "       W_r^{(l)} = \\sum_{b=1}^B a_{rb}^{(l)}V_b^{(l)}\n",
    " \\end{equation}\n",
    " \n",
    "where : <br/>\n",
    "$B$ is the number of bases\n",
    "$V_b^{(l)}$ are linearly combined with coefficients $a_{rb}^{(l)}$.\n",
    "\n",
    "Therefore, the weight $W_r^{(l)}$ is a linear combination of basis transformation V(l)b with coefficients a(l)rb. The number of bases B is much smaller than the number of relations in the knowledge base.\n",
    "\n",
    "\n",
    "#### Decoder :  Link prediction\n",
    "\n",
    "use DistMult factorization (Yang et al. 2014) as the scoring function, which is known to perform well on standard link prediction benchmarks when used on its own:\n",
    "\n",
    "\\begin{equation}\n",
    "       f(s, r, o) = e^T_sR_re_o . \n",
    " \\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "082e9718",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "from dputils import load_data, calc_mrr, evaluate\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn.conv import MessagePassing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7400cc4",
   "metadata": {},
   "source": [
    "#### Load WN18 Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7d731549",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load data from ./data/wn18\n",
      "num_entity: 40943\n",
      "num_relation: 18\n",
      "num_train_triples: 141442\n",
      "num_valid_triples: 5000\n",
      "num_test_triples: 5000\n"
     ]
    }
   ],
   "source": [
    "entity2id, relation2id, train_triplets, valid_triplets, test_triplets = load_data('./data/wn18')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "56c8058e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "numpy.ndarray"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(train_triplets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "86d00767",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 8881,    12, 32231],\n",
       "       [35365,    17, 15368],\n",
       "       [13944,    16, 24784],\n",
       "       [35723,     8,  3497],\n",
       "       [36279,     8, 21979]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_triplets[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e000eae1",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_entity = len(entity2id)\n",
    "num_relation = len(relation2id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ca0a440",
   "metadata": {},
   "source": [
    "#### Prep Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8603c6e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def negative_sampling(pos_samples, num_entity, negative_rate):\n",
    "    \n",
    "    size_of_batch = len(pos_samples)\n",
    "    num_to_generate = size_of_batch * negative_rate\n",
    "    neg_samples = np.tile(pos_samples, (negative_rate, 1))\n",
    "    labels = np.zeros(size_of_batch * (negative_rate + 1), dtype=np.float32)\n",
    "    labels[: size_of_batch] = 1\n",
    "    values = np.random.choice(num_entity, size=num_to_generate)\n",
    "    choices = np.random.uniform(size=num_to_generate)\n",
    "    subj = choices > 0.5\n",
    "    obj = choices <= 0.5\n",
    "    neg_samples[subj, 0] = values[subj]\n",
    "    neg_samples[obj, 2] = values[obj]\n",
    "\n",
    "    return np.concatenate((pos_samples, neg_samples)), labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a6a0e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_scatter import scatter_add\n",
    "def edge_normalization(edge_type, edge_index, num_entity, num_relation):\n",
    "    '''\n",
    "        Edge normalization trick\n",
    "        - one_hot: (num_edge, num_relation)\n",
    "        - deg: (num_node, num_relation)\n",
    "        - index: (num_edge)\n",
    "        - deg[edge_index[0]]: (num_edge, num_relation)\n",
    "        - edge_norm: (num_edge)\n",
    "    '''\n",
    "    one_hot = F.one_hot(edge_type.long(), num_classes = 2 * num_relation).to(torch.float)\n",
    "    deg = scatter_add(one_hot, edge_index[0], dim = 0, dim_size = num_entity)\n",
    "    index = edge_type + torch.arange(len(edge_index[0])) * (2 * num_relation)\n",
    "    edge_norm = 1 / deg[edge_index[0]].view(-1)[index]\n",
    "\n",
    "    return edge_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cceafc08",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_sampled_graph_and_labels(triplets, sample_size, split_size, num_entity, num_relation, negative_rate):\n",
    "    \"\"\"\n",
    "        Get training graph and signals\n",
    "        First perform edge neighborhood sampling on graph, then perform negative\n",
    "        sampling to generate negative samples\n",
    "    \"\"\"\n",
    "    def sample_edge_uniform(n_triples, sample_size):\n",
    "        \"\"\"Sample edges uniformly from all the edges.\"\"\"\n",
    "        all_edges = np.arange(n_triples)\n",
    "        return np.random.choice(all_edges, sample_size, replace=False)\n",
    "\n",
    "    edges = sample_edge_uniform(len(triplets), sample_size)\n",
    "\n",
    "    # Select sampled edges\n",
    "    edges = triplets[edges]\n",
    "    src, rel, dst = edges.transpose()\n",
    "    uniq_entity, edges = np.unique((src, dst), return_inverse=True)\n",
    "    src, dst = np.reshape(edges, (2, -1))\n",
    "    relabeled_edges = np.stack((src, rel, dst)).transpose()\n",
    "\n",
    "    # Negative sampling\n",
    "    samples, labels = negative_sampling(relabeled_edges, len(uniq_entity), negative_rate)\n",
    "\n",
    "    # further split graph, only half of the edges will be used as graph in message passing\n",
    "    # structure, while the rest half is used as unseen positive samples\n",
    "    split_size = int(sample_size * split_size)\n",
    "    graph_split_ids = np.random.choice(np.arange(sample_size),\n",
    "                                       size=split_size, replace=False)\n",
    "\n",
    "    src = torch.tensor(src[graph_split_ids], dtype = torch.long).contiguous()\n",
    "    dst = torch.tensor(dst[graph_split_ids], dtype = torch.long).contiguous()\n",
    "    rel = torch.tensor(rel[graph_split_ids], dtype = torch.long).contiguous()\n",
    "\n",
    "    # Create bi-directional graph\n",
    "    src, dst = torch.cat((src, dst)), torch.cat((dst, src))\n",
    "    rel = torch.cat((rel, rel + num_relation))\n",
    "\n",
    "    edge_index = torch.stack((src, dst))\n",
    "    edge_type = rel\n",
    "\n",
    "\n",
    "    entity = torch.from_numpy(uniq_entity)\n",
    "    edge_norm = edge_normalization(edge_type, edge_index, len(uniq_entity), num_relation)\n",
    "    samples = torch.from_numpy(samples)\n",
    "    labels = torch.from_numpy(labels)\n",
    "    \n",
    "    '''\n",
    "    # use in message propogation\n",
    "    entity : sampled unique entity ids, [0 < x < 2N]\n",
    "    edge_index: [2, N], pairs of (src, dst) node ids, positive only, bi-directional\n",
    "    edge_type: [N], bi-directional 0<x<36\n",
    "    edge_norm: [N], normalizating factor\n",
    "    \n",
    "    # use in loss calculation (single_direction)\n",
    "    samples: [3, 2N] : first half pos, sec half neg\n",
    "    labels: [2N]: [1,1,1,...0,0,0 ]\n",
    "    '''\n",
    "    \n",
    "    return entity, edge_index, edge_type, edge_norm, samples, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8123767",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_eval_graph(num_entity, num_relation, triplets):\n",
    "    src, rel, dst = triplets.transpose()\n",
    "\n",
    "    src = torch.from_numpy(src)\n",
    "    rel = torch.from_numpy(rel)\n",
    "    dst = torch.from_numpy(dst)\n",
    "\n",
    "    src, dst = torch.cat((src, dst)), torch.cat((dst, src))\n",
    "    rel = torch.cat((rel, rel + num_relation))\n",
    "\n",
    "    edge_index = torch.stack((src, dst)).long()\n",
    "    edge_type = rel\n",
    "\n",
    "    entity = torch.from_numpy(np.arange(num_entity))\n",
    "    edge_norm = edge_normalization(edge_type, edge_index, num_entity, num_relation)\n",
    "\n",
    "    return entity, edge_index, edge_type, edge_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "005939ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "#full graph without sampling/neg sampling and spliting\n",
    "eval_graph = build_eval_graph(num_entity, num_relation, train_triplets) \n",
    "\n",
    "all_triplets = torch.LongTensor(np.concatenate((train_triplets, valid_triplets, test_triplets)))\n",
    "valid_triplets = torch.LongTensor(valid_triplets)\n",
    "test_triplets = torch.LongTensor(test_triplets)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0ee42d1",
   "metadata": {},
   "source": [
    "#### RGCN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bee1dfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RGCNConv(MessagePassing):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        in_channels (int): Size of each input sample.\n",
    "        out_channels (int): Size of each output sample.\n",
    "        num_relations (int): Number of relations.\n",
    "        num_bases (int): Number of bases used for basis-decomposition.\n",
    "        root_weight (bool, optional): If set to :obj:`False`, the layer will\n",
    "            not add transformed root node features to the output.\n",
    "            (default: :obj:`True`)\n",
    "        bias (bool, optional): If set to :obj:`False`, the layer will not learn\n",
    "            an additive bias. (default: :obj:`True`)\n",
    "        **kwargs (optional): Additional arguments of\n",
    "            :class:`torch_geometric.nn.conv.MessagePassing`.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, num_relations, num_bases,\n",
    "                 root_weight=True, bias=True, **kwargs):\n",
    "        super(RGCNConv, self).__init__(aggr='mean', **kwargs)\n",
    "\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.num_relations = num_relations\n",
    "        self.num_bases = num_bases\n",
    "\n",
    "        self.basis = nn.Parameter(torch.Tensor(num_bases, in_channels, out_channels))\n",
    "        self.att = nn.Parameter(torch.Tensor(num_relations, num_bases))\n",
    "\n",
    "        if root_weight:\n",
    "            self.root = nn.Parameter(torch.Tensor(in_channels, out_channels))\n",
    "        else:\n",
    "            self.register_parameter('root', None)\n",
    "\n",
    "        if bias:\n",
    "            self.bias = nn.Parameter(torch.Tensor(out_channels))\n",
    "        else:\n",
    "            self.register_parameter('bias', None)\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        size = self.num_bases * self.in_channels\n",
    "        uniform(size, self.basis)\n",
    "        uniform(size, self.att)\n",
    "        uniform(size, self.root)\n",
    "        uniform(size, self.bias)\n",
    "\n",
    "\n",
    "    def forward(self, x, edge_index, edge_type, edge_norm=None, size=None):\n",
    "        \"\"\"\"\"\"\n",
    "        return self.propagate(edge_index, size=size, x=x, edge_type=edge_type,\n",
    "                              edge_norm=edge_norm)\n",
    "\n",
    "\n",
    "    def message(self, x_j, edge_index_j, edge_type, edge_norm):\n",
    "        w = torch.matmul(self.att, self.basis.view(self.num_bases, -1))\n",
    "\n",
    "        # If no node features are given, we implement a simple embedding\n",
    "        # loopkup based on the target node index and its edge type.\n",
    "        if x_j is None:\n",
    "            w = w.view(-1, self.out_channels)\n",
    "            index = edge_type * self.in_channels + edge_index_j\n",
    "            out = torch.index_select(w, 0, index)\n",
    "        else:\n",
    "            w = w.view(self.num_relations, self.in_channels, self.out_channels)\n",
    "            w = torch.index_select(w, 0, edge_type)\n",
    "            out = torch.bmm(x_j.unsqueeze(1), w).squeeze(-2)\n",
    "\n",
    "        return out if edge_norm is None else out * edge_norm.view(-1, 1)\n",
    "\n",
    "    def update(self, aggr_out, x):\n",
    "        if self.root is not None:\n",
    "            if x is None:\n",
    "                out = aggr_out + self.root\n",
    "            else:\n",
    "                out = aggr_out + torch.matmul(x, self.root)\n",
    "\n",
    "        if self.bias is not None:\n",
    "            out = out + self.bias\n",
    "        return out\n",
    "\n",
    "    def __repr__(self):\n",
    "        return '{}({}, {}, num_relations={})'.format(\n",
    "            self.__class__.__name__, self.in_channels, self.out_channels,\n",
    "            self.num_relations)\n",
    "\n",
    "def uniform(size, tensor):\n",
    "    bound = 1.0 / math.sqrt(size)\n",
    "    if tensor is not None:\n",
    "        tensor.data.uniform_(-bound, bound)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b11c6d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RGCN(torch.nn.Module):\n",
    "    def __init__(self, num_entities, num_relations, num_bases, dropout):\n",
    "        super(RGCN, self).__init__()\n",
    "\n",
    "        self.entity_embedding = nn.Embedding(num_entities, 100)\n",
    "        self.relation_embedding = nn.Parameter(torch.Tensor(num_relations, 100))\n",
    "\n",
    "        nn.init.xavier_uniform_(self.relation_embedding, gain=nn.init.calculate_gain('relu'))\n",
    "\n",
    "        self.conv1 = RGCNConv(\n",
    "            100, 100, num_relations * 2, num_bases=num_bases)\n",
    "        self.conv2 = RGCNConv(\n",
    "            100, 100, num_relations * 2, num_bases=num_bases)\n",
    "\n",
    "        self.dropout_ratio = dropout\n",
    "\n",
    "    def forward(self, entity, edge_index, edge_type, edge_norm):\n",
    "        x = self.entity_embedding(entity)\n",
    "        x = F.relu(self.conv1(x, edge_index, edge_type, edge_norm))\n",
    "        x = F.dropout(x, p = self.dropout_ratio, training = self.training)\n",
    "        x = self.conv2(x, edge_index, edge_type, edge_norm)\n",
    "        \n",
    "        return x\n",
    "\n",
    "    def distmult(self, embedding, triplets):\n",
    "        s = embedding[triplets[:,0]]\n",
    "        r = self.relation_embedding[triplets[:,1]]\n",
    "        o = embedding[triplets[:,2]]\n",
    "        score = torch.sum(s * r * o, dim=1)\n",
    "        \n",
    "        return score\n",
    "\n",
    "    def score_loss(self, embedding, triplets, target):\n",
    "        score = self.distmult(embedding, triplets)\n",
    "\n",
    "        return F.binary_cross_entropy_with_logits(score, target)\n",
    "\n",
    "    def reg_loss(self, embedding):\n",
    "        return torch.mean(embedding.pow(2)) + torch.mean(self.relation_embedding.pow(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fe5270d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_size = 3000\n",
    "split_size = 0.5\n",
    "negative_rate = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4b28c5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 10000\n",
    "batch_size = 5000\n",
    "evaluate_interval = 1000\n",
    "num_bases = 4\n",
    "dropout = 0.2\n",
    "lr = 0.01 \n",
    "grad_norm = 1.0\n",
    "regularization = 1e-2\n",
    "cuda = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fb4c4d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = RGCN(num_entity, num_relation, num_bases=num_bases, dropout=dropout)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr = lr )\n",
    "print(model)\n",
    "if cuda:\n",
    "    model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b74cf22",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "entity, edge_index, edge_type, edge_norm, samples, labels = generate_sampled_graph_and_labels(train_triplets,\n",
    "                                                                                  sample_size = batch_size,\n",
    "                                                                                  split_size = 0.5,\n",
    "                                                                                  num_entity = num_entity,\n",
    "                                                                                  num_relation = num_relation,\n",
    "                                                                                  negative_rate = negative_rate)\n",
    "\n",
    "if cuda:\n",
    "    entity = entity.cuda()\n",
    "    edge_index = edge_index.cuda()\n",
    "    edge_type = edge_type.cuda()\n",
    "    edge_norm = edge_norm.cuda()\n",
    "    samples = samples.cuda()\n",
    "    labels = labels.cuda()\n",
    "\n",
    "entity_embedding = model(entity, edge_index, edge_type, edge_norm)\n",
    "loss = model.score_loss(entity_embedding, samples, labels) + regularization * model.reg_loss(entity_embedding)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "932d9667",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "best_mrr  = 0.\n",
    "for epoch in range(1, (num_epochs + 1)):\n",
    "\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # sample training data\n",
    "    entity, edge_index, edge_type, edge_norm, samples, labels = generate_sampled_graph_and_labels(train_triplets,\n",
    "                                                                                      sample_size = batch_size,\n",
    "                                                                                      split_size = 0.5,\n",
    "                                                                                      num_entity = num_entity,\n",
    "                                                                                      num_relation = num_relation,\n",
    "                                                                                      negative_rate = negative_rate)\n",
    "\n",
    "    if cuda:\n",
    "        entity = entity.cuda()\n",
    "        edge_index = edge_index.cuda()\n",
    "        edge_type = edge_type.cuda()\n",
    "        edge_norm = edge_norm.cuda()\n",
    "        samples = samples.cuda()\n",
    "        labels = labels.cuda()\n",
    "    \n",
    "    # traning and back prop\n",
    "    entity_embedding = model(entity, edge_index, edge_type, edge_norm)\n",
    "    loss = model.score_loss(entity_embedding, samples, labels) + regularization * model.reg_loss(entity_embedding)\n",
    "    loss.backward()\n",
    "    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)\n",
    "    optimizer.step()\n",
    "\n",
    "    if epoch % evaluate_interval == 0:\n",
    "\n",
    "        print(\"Train Loss {} at epoch {}\".format(loss, epoch))\n",
    "\n",
    "#         if cuda:\n",
    "#             model.cpu()\n",
    "\n",
    "#         model.eval()\n",
    "#         valid_mrr = evaluate(valid_triplets, model, eval_graph, all_triplets)\n",
    "\n",
    "#         if valid_mrr > best_mrr:\n",
    "#             best_mrr = valid_mrr\n",
    "#             torch.save({'state_dict': model.state_dict(), 'epoch': epoch},\n",
    "#                         'best_mrr_model.pth')\n",
    "\n",
    "#         if cuda:\n",
    "#             model.cuda()\n",
    "\n",
    "if cuda:\n",
    "    model.cpu()\n",
    "\n",
    "model.eval()\n",
    "\n",
    "# checkpoint = torch.load('best_mrr_model.pth')\n",
    "# model.load_state_dict(checkpoint['state_dict'])\n",
    "# test_mrr = evaluate(test_triplets, model, eval_graph, all_triplets)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "801c248c",
   "metadata": {},
   "source": [
    "<!-- Train Loss 0.13028721511363983 at epoch 1000\n",
    "Train Loss 0.09021419286727905 at epoch 2000\n",
    "Train Loss 0.07641539722681046 at epoch 3000\n",
    "Train Loss 0.06500016152858734 at epoch 4000\n",
    "Train Loss 0.059444960206747055 at epoch 5000\n",
    "Train Loss 0.056531891226768494 at epoch 6000\n",
    "Train Loss 0.052266258746385574 at epoch 7000\n",
    "Train Loss 0.04904014989733696 at epoch 8000\n",
    "Train Loss 0.05327734351158142 at epoch 9000\n",
    "Train Loss 0.03978230431675911 at epoch 10000\n",
    "\n",
    "MRR (filtered): 0.191787\n",
    "Hits (filtered) @ 1: 0.116500\n",
    "Hits (filtered) @ 3: 0.206700\n",
    "Hits (filtered) @ 10: 0.340600 -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a92cffa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "Train Loss 0.13028721511363983 at epoch 1000\n",
    "Train Loss 0.09021419286727905 at epoch 2000\n",
    "Train Loss 0.07641539722681046 at epoch 3000\n",
    "Train Loss 0.06500016152858734 at epoch 4000\n",
    "Train Loss 0.059444960206747055 at epoch 5000\n",
    "Train Loss 0.056531891226768494 at epoch 6000\n",
    "Train Loss 0.052266258746385574 at epoch 7000\n",
    "Train Loss 0.04904014989733696 at epoch 8000\n",
    "Train Loss 0.05327734351158142 at epoch 9000\n",
    "Train Loss 0.03978230431675911 at epoch 10000\n",
    "\n",
    "MRR (filtered): 0.191787\n",
    "Hits (filtered) @ 1: 0.116500\n",
    "Hits (filtered) @ 3: 0.206700\n",
    "Hits (filtered) @ 10: 0.340600"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pytorch-gpu] *",
   "language": "python",
   "name": "conda-env-pytorch-gpu-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
